import argparse


def args_parser():
    parser = argparse.ArgumentParser()
    # python main.py --global_model=CNN2 -local_model=CNN2 --dataset=cifar10 --E=5 --R=50
    # Total
    parser.add_argument('--device', type=str, default='cuda:0',
                        help="device: {cuda, cpu}")
    parser.add_argument('--num', type=int, default=1,
                        help="Number of exp")
    parser.add_argument('--node_num', type=int, default=5,
                        help="Number of nodes")
    parser.add_argument('--R', type=int, default=50,
                        help="Number of rounds: R")
    parser.add_argument('--E', type=int, default=5,
                        help="Number of local epochs: E")
    parser.add_argument('--frac', type=float, default=1,  # 选择比例
                        help='the fraction of clients: C')
    parser.add_argument('--iid', type=int, default=1,  # 选择比例
                        help='set 1 for iid')
    parser.add_argument('--dataset', type=str, default='cifar10',
                        help="Type of algorithms:{mnist, cifar10,cifar100}")

    # Model
    parser.add_argument('--global_model', type=str, default='CNN1',
                        help='Type of global model: {mlp, LeNet5,CNN1,CNN2 ResNet18}')
    parser.add_argument('--local_model', type=str, default='CNN_fedmd',
                        help='Type of local model: {mlp, LeNet5,CNN, ResNet18}')

    # Data
    parser.add_argument('--batchsize', type=int, default=128,
                        help="batchsize")

    # Optim
    parser.add_argument('--optimizer', type=str, default='sgd',
                        help="optimizer: {sgd, adam}")
    parser.add_argument('--lr', type=float, default=0.1,
                        help='learning rate')
    parser.add_argument('--lr_step', type=int, default=10,
                        help='learning rate decay step size')
    parser.add_argument('--stop_decay', type=int, default=40,
                        help='round when learning rate stop decay')
    parser.add_argument('--momentum', type=float, default=0.9,
                        help='SGD momentum')

    args = parser.parse_args()
    return args
